"""
搜索最优的 K 值并将最优 KNN 模型保存到磁盘。

行为：
 - 加载手写数字数据集（sklearn 的 digits）
 - 将数据标准化并划分为训练/测试集
 - 在 k=1..20 上训练 KNeighborsClassifier，记录测试集准确率
 - 使用 tqdm 显示进度条
 - 保存最佳模型为 pickle 文件
 - 打印最优准确率和对应的 k

输出文件：best_knn_model.pkl（保存在脚本相同目录下）
"""

import os
import pickle
from sklearn.datasets import fetch_openml, load_digits
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.neighbors import KNeighborsClassifier
from sklearn.metrics import accuracy_score
from tqdm import tqdm
from sklearn.utils import shuffle
import numpy as np
from sklearn.svm import SVC

def augment_data(X, y):
    """对数据进行增强，生成旋转、缩放等变体"""
    augmented_X, augmented_y = [], []
    for img, label in zip(X, y):
        img = img.reshape(8, 8)  # 恢复为二维图像
        augmented_X.append(img)
        augmented_y.append(label)

        # 数据增强：旋转
        rotated = np.rot90(img)
        augmented_X.append(rotated)
        augmented_y.append(label)

        # 数据增强：缩放
        scaled = img * 1.2  # 简单放大
        scaled = np.clip(scaled, 0, 16)  # 限制范围
        augmented_X.append(scaled)
        augmented_y.append(label)

    augmented_X = np.array(augmented_X).reshape(-1, 64)  # 扁平化
    augmented_y = np.array(augmented_y)
    return shuffle(augmented_X, augmented_y)

def find_and_save_best_knn(output_path: str = "best_knn_model.pkl", k_min: int = 1, k_max: int = 20, test_size: float = 0.2, random_state: int = 42):
	"""在 k_min..k_max 搜索最优 KNN，并将最佳模型保存为 pickle。"""

	# 加载 sklearn digits 数据集
	print("Loading sklearn digits dataset...")
	digits = load_digits()
	X, y = digits.data, digits.target

	# 数据增强
	print("Augmenting data...")
	X, y = augment_data(X, y)

	# 划分数据集
	X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=test_size, random_state=random_state, stratify=y)

	# 标准化
	scaler = StandardScaler()
	X_train = scaler.fit_transform(X_train)
	X_test = scaler.transform(X_test)

	best_k = None
	best_acc = 0.0
	best_model = None
	acc_list = []

	# 搜索 k
	for k in tqdm(range(k_min, k_max + 1), desc="Searching k", unit="k"):
		knn = KNeighborsClassifier(n_neighbors=k)
		knn.fit(X_train, y_train)
		preds = knn.predict(X_test)
		acc = accuracy_score(y_test, preds)
		acc_list.append(acc)

		if acc > best_acc:
			best_acc = acc
			best_k = k
			# 保存包含 scaler 的元组，以便在推理时也能复用相同的标准化器
			best_model = (scaler, knn)

	# 将最佳模型保存到文件
	save_dir = os.path.dirname(os.path.abspath(output_path)) or os.getcwd()
	os.makedirs(save_dir, exist_ok=True)
	with open(output_path, "wb") as f:
		pickle.dump(best_model, f)
	print(f"Saved best model to: {os.path.abspath(output_path)}")

	# 绘制准确率随 k 变化的图
	try:
		import matplotlib.pyplot as plt

		ks = list(range(k_min, k_max + 1))
		plt.figure(figsize=(8, 5))
		plt.plot(ks, acc_list, marker='o', linestyle='-')
		plt.xlabel('k')
		plt.ylabel('Accuracy')
		plt.title('K vs Accuracy')
		plt.axvline(x=best_k, color='red', linestyle='--', label=f'Best k = {best_k}')
		plt.scatter([best_k], [best_acc], color='red')
		plt.text(best_k + 0.5, best_acc, f'k={best_k}\n{best_acc*100:.2f}%', color='red')
		plt.grid(True)

		plot_path = os.path.join(save_dir, 'accuracy_plot.pdf')
		plt.tight_layout()
		plt.savefig(plot_path, format='pdf')
		plt.close()
		print(f"Saved accuracy plot to: {os.path.abspath(plot_path)}")
	except Exception as e:
		print(f"Error while plotting: {e}")

	print(f"Best k: {best_k}")
	print(f"Best accuracy: {best_acc * 100:.2f}%")

	return best_k, best_acc, acc_list

def find_and_save_best_model(output_path: str = "best_model.pkl", test_size: float = 0.2, random_state: int = 42):
    """训练支持向量机（SVM）模型，并将最佳模型保存为 pickle。"""

    # 加载 sklearn digits 数据集
    print("Loading sklearn digits dataset...")
    digits = load_digits()
    X, y = digits.data, digits.target

    # TODO: 加入用户手写样本（需要用户提供样本数据）

    # 划分数据集
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=test_size, random_state=random_state, stratify=y)

    # 标准化
    scaler = StandardScaler()
    X_train = scaler.fit_transform(X_train)
    X_test = scaler.transform(X_test)

    # 训练 SVM 模型
    print("Training SVM model...")
    svm = SVC(kernel='rbf', C=10, gamma=0.01)
    svm.fit(X_train, y_train)

    # 测试模型
    acc = svm.score(X_test, y_test)
    print(f"Model accuracy: {acc * 100:.2f}%")

    # 保存模型
    save_dir = os.path.dirname(os.path.abspath(output_path)) or os.getcwd()
    os.makedirs(save_dir, exist_ok=True)
    with open(output_path, "wb") as f:
        pickle.dump((scaler, svm), f)
    print(f"Saved model to: {os.path.abspath(output_path)}")

if __name__ == "__main__":
	# 默认在当前目录保存
	try:
		print('Script start', flush=True)
		out_file = os.path.join(os.path.dirname(__file__), "best_knn_model.pkl")
		print('Output path will be:', out_file, flush=True)
		best_k, best_acc, acc_list = find_and_save_best_knn(output_path=out_file)
		print('Finished successfully: best_k=', best_k, 'best_acc=', best_acc, flush=True)
	except Exception as exc:
		# 显式打印 traceback，确保重定向到 run_out.log 时也能捕获详细错误
		import traceback
		print('Script failed with exception:', flush=True)
		traceback.print_exc()
